#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
TCP迁移方案综合测试

测试内容：
1. 内存缓存系统
2. HTTP API接口
3. 数据上传和存储
4. 实时数据获取
5. 版本控制机制
6. 内存管理和清理
7. 并发安全性
8. 性能监控
"""

import asyncio
import json
import time
import random
import threading
from datetime import datetime, timezone
from typing import Dict, List, Optional
from collections import deque
from threading import RLock
import httpx
import pytest
from unittest.mock import Mock, patch

# 模拟内存缓存系统
class MemoryCache:
    """内存缓存系统实现"""
    
    def __init__(self, max_rooms=100, max_window_size=300, cache_ttl=3600):
        self._lock = RLock()
        
        # 房间最新数据缓存 {room_id: latest_data}
        self.latest_data: Dict[str, Dict] = {}
        
        # 房间数据版本号 {room_id: version}
        self.data_versions: Dict[str, int] = {}
        
        # 时间窗口数据缓存 {room_id: deque([(timestamp, data), ...])}
        self.time_window_data: Dict[str, deque] = {}
        
        # 缓存过期时间 {room_id: expire_timestamp}
        self.cache_expiry: Dict[str, float] = {}
        
        # 配置参数
        self.max_rooms = max_rooms
        self.max_window_size = max_window_size
        self.cache_ttl = cache_ttl
        
        # 性能监控
        self.metrics = {
            'cache_hits': 0,
            'cache_misses': 0,
            'data_updates': 0,
            'cleanup_runs': 0
        }
    
    def add_data(self, room_id: str, data: Dict) -> None:
        """添加数据到缓存"""
        with self._lock:
            current_time = time.time()
            
            # 限制房间数量
            if len(self.latest_data) >= self.max_rooms and room_id not in self.latest_data:
                self._evict_oldest_room()
            
            # 更新最新数据
            self.latest_data[room_id] = data
            
            # 更新版本号
            self.data_versions[room_id] = self.data_versions.get(room_id, 0) + 1
            
            # 添加到时间窗口
            if room_id not in self.time_window_data:
                self.time_window_data[room_id] = deque()
            
            self.time_window_data[room_id].append((current_time, data))
            
            # 限制时间窗口大小
            while len(self.time_window_data[room_id]) > self.max_window_size:
                self.time_window_data[room_id].popleft()
            
            # 更新过期时间
            self.cache_expiry[room_id] = current_time + self.cache_ttl
            
            self.metrics['data_updates'] += 1
    
    def get_latest_data(self, room_id: str) -> Optional[Dict]:
        """获取最新数据"""
        with self._lock:
            if room_id in self.latest_data:
                self.metrics['cache_hits'] += 1
                return self.latest_data[room_id]
            else:
                self.metrics['cache_misses'] += 1
                return None
    
    def get_data_since_timestamp(self, room_id: str, since_timestamp: float) -> List[Dict]:
        """获取指定时间戳之后的数据"""
        with self._lock:
            if room_id not in self.time_window_data:
                self.metrics['cache_misses'] += 1
                return []
            
            recent_data = []
            for timestamp, data in self.time_window_data[room_id]:
                if timestamp > since_timestamp:
                    recent_data.append(data)
            
            if recent_data:
                self.metrics['cache_hits'] += 1
            else:
                self.metrics['cache_misses'] += 1
            
            return recent_data
    
    def get_version(self, room_id: str) -> int:
        """获取房间数据版本号"""
        with self._lock:
            return self.data_versions.get(room_id, 0)
    
    def check_version_update(self, room_id: str, client_version: int) -> bool:
        """检查是否有新版本数据"""
        current_version = self.get_version(room_id)
        return current_version > client_version
    
    def cleanup_expired_data(self) -> int:
        """清理过期数据"""
        with self._lock:
            current_time = time.time()
            cleaned_count = 0
            
            # 清理过期的缓存
            expired_rooms = []
            for room_id, expire_time in self.cache_expiry.items():
                if current_time > expire_time:
                    expired_rooms.append(room_id)
            
            for room_id in expired_rooms:
                self._remove_room_cache(room_id)
                cleaned_count += 1
            
            # 清理时间窗口数据（保留最近5分钟）
            window_seconds = 300
            cutoff_time = current_time - window_seconds
            
            for room_id, data_queue in self.time_window_data.items():
                while data_queue and data_queue[0][0] < cutoff_time:
                    data_queue.popleft()
                    cleaned_count += 1
            
            self.metrics['cleanup_runs'] += 1
            return cleaned_count
    
    def _evict_oldest_room(self) -> None:
        """驱逐最旧的房间数据"""
        if not self.cache_expiry:
            return
        
        oldest_room = min(self.cache_expiry.keys(), 
                         key=lambda x: self.cache_expiry[x])
        self._remove_room_cache(oldest_room)
    
    def _remove_room_cache(self, room_id: str) -> None:
        """移除房间缓存"""
        self.latest_data.pop(room_id, None)
        self.data_versions.pop(room_id, None)
        self.time_window_data.pop(room_id, None)
        self.cache_expiry.pop(room_id, None)
    
    def get_cache_stats(self) -> Dict:
        """获取缓存统计信息"""
        with self._lock:
            total_requests = self.metrics['cache_hits'] + self.metrics['cache_misses']
            hit_rate = self.metrics['cache_hits'] / total_requests if total_requests > 0 else 0
            
            return {
                'total_rooms': len(self.latest_data),
                'total_requests': total_requests,
                'cache_hit_rate': hit_rate,
                'data_updates': self.metrics['data_updates'],
                'cleanup_runs': self.metrics['cleanup_runs']
            }


# 模拟EEG数据生成器
class EEGDataGenerator:
    """EEG数据生成器"""
    
    def __init__(self):
        self.states = ['awake', 'light_sleep', 'deep_sleep', 'rem_sleep']
        self.sequence = 0
    
    def generate_eeg_data(self, room_id: str, session_id: str) -> Dict:
        """生成模拟EEG数据"""
        self.sequence += 1
        
        return {
            'room_id': room_id,
            'session_id': session_id,
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'sequence': self.sequence,
            'data': {
                'state': random.choice(self.states),
                'confidence': round(random.uniform(0.7, 0.95), 2),
                'brainwave_ratios': {
                    'delta': round(random.uniform(0.2, 0.4), 2),
                    'theta': round(random.uniform(0.15, 0.35), 2),
                    'alpha': round(random.uniform(0.1, 0.3), 2),
                    'beta': round(random.uniform(0.1, 0.25), 2),
                    'gamma': round(random.uniform(0.05, 0.15), 2)
                },
                'raw_channels': {
                    f'channel_{i}': [round(random.uniform(-50, 50), 2) for _ in range(10)]
                    for i in range(1, 9)
                },
                'signal_quality': round(random.uniform(0.8, 0.98), 2),
                'battery_level': random.randint(20, 100)
            }
        }


# 模拟HTTP客户端
class MockHTTPClient:
    """模拟HTTP客户端"""
    
    def __init__(self, memory_cache: MemoryCache):
        self.memory_cache = memory_cache
        self.base_url = "http://localhost:8000"
    
    async def upload_data(self, data: Dict) -> Dict:
        """模拟数据上传"""
        # 模拟网络延迟
        await asyncio.sleep(0.01)
        
        # 添加到内存缓存
        self.memory_cache.add_data(data['room_id'], data)
        
        return {
            'status': 'success',
            'timestamp': datetime.now(timezone.utc).isoformat(),
            'version': self.memory_cache.get_version(data['room_id'])
        }
    
    async def get_realtime_data(self, room_id: str, since: Optional[float] = None) -> Dict:
        """获取实时数据"""
        await asyncio.sleep(0.005)
        
        latest = self.memory_cache.get_latest_data(room_id)
        recent = []
        
        if since is not None:
            recent = self.memory_cache.get_data_since_timestamp(room_id, since)
        
        return {
            'room_id': room_id,
            'version': self.memory_cache.get_version(room_id),
            'last_update': datetime.now(timezone.utc).isoformat(),
            'data': {
                'latest': latest,
                'recent': recent
            }
        }
    
    async def check_version(self, room_id: str) -> Dict:
        """检查版本"""
        await asyncio.sleep(0.002)
        
        version = self.memory_cache.get_version(room_id)
        latest_data = self.memory_cache.get_latest_data(room_id)
        
        return {
            'room_id': room_id,
            'version': version,
            'last_update': datetime.now(timezone.utc).isoformat(),
            'has_new_data': latest_data is not None
        }


# 测试类
class TestTCPMigration:
    """TCP迁移综合测试"""
    
    def setup_method(self):
        """测试前准备"""
        self.memory_cache = MemoryCache()
        self.data_generator = EEGDataGenerator()
        self.http_client = MockHTTPClient(self.memory_cache)
    
    def test_memory_cache_basic_operations(self):
        """测试内存缓存基本操作"""
        room_id = "room_test_01"
        test_data = self.data_generator.generate_eeg_data(room_id, "session_001")
        
        # 测试添加数据
        self.memory_cache.add_data(room_id, test_data)
        
        # 测试获取最新数据
        latest = self.memory_cache.get_latest_data(room_id)
        assert latest is not None
        assert latest['room_id'] == room_id
        
        # 测试版本号
        version = self.memory_cache.get_version(room_id)
        assert version == 1
        
        # 添加更多数据
        for i in range(5):
            new_data = self.data_generator.generate_eeg_data(room_id, "session_001")
            self.memory_cache.add_data(room_id, new_data)
        
        # 检查版本号递增
        assert self.memory_cache.get_version(room_id) == 6
    
    def test_time_window_data(self):
        """测试时间窗口数据"""
        room_id = "room_test_02"
        
        # 添加多条数据
        timestamps = []
        for i in range(10):
            data = self.data_generator.generate_eeg_data(room_id, "session_002")
            self.memory_cache.add_data(room_id, data)
            timestamps.append(time.time())
            time.sleep(0.01)  # 确保时间戳不同
        
        # 测试获取指定时间后的数据
        mid_timestamp = timestamps[5]
        recent_data = self.memory_cache.get_data_since_timestamp(room_id, mid_timestamp)
        
        # 应该获取到后5条数据
        assert len(recent_data) >= 4  # 考虑时间精度问题
    
    def test_cache_expiry_and_cleanup(self):
        """测试缓存过期和清理"""
        # 创建短TTL的缓存
        short_cache = MemoryCache(cache_ttl=0.1)  # 0.1秒过期
        
        room_id = "room_test_03"
        data = self.data_generator.generate_eeg_data(room_id, "session_003")
        
        # 添加数据
        short_cache.add_data(room_id, data)
        assert short_cache.get_latest_data(room_id) is not None
        
        # 等待过期
        time.sleep(0.2)
        
        # 执行清理
        cleaned = short_cache.cleanup_expired_data()
        assert cleaned > 0
        
        # 数据应该被清理
        assert short_cache.get_latest_data(room_id) is None
    
    def test_concurrent_access(self):
        """测试并发访问安全性"""
        room_id = "room_test_04"
        results = []
        
        def worker(worker_id):
            for i in range(50):
                data = self.data_generator.generate_eeg_data(room_id, f"session_{worker_id}")
                self.memory_cache.add_data(room_id, data)
                
                latest = self.memory_cache.get_latest_data(room_id)
                if latest:
                    results.append(latest)
        
        # 启动多个线程
        threads = []
        for i in range(5):
            thread = threading.Thread(target=worker, args=(i,))
            threads.append(thread)
            thread.start()
        
        # 等待所有线程完成
        for thread in threads:
            thread.join()
        
        # 检查结果
        assert len(results) > 0
        final_version = self.memory_cache.get_version(room_id)
        assert final_version == 250  # 5个线程 * 50次操作
    
    @pytest.mark.asyncio
    async def test_http_upload_flow(self):
        """测试HTTP上传流程"""
        room_id = "room_test_05"
        session_id = "session_005"
        
        # 生成测试数据
        test_data = self.data_generator.generate_eeg_data(room_id, session_id)
        
        # 上传数据
        response = await self.http_client.upload_data(test_data)
        
        assert response['status'] == 'success'
        assert 'timestamp' in response
        assert response['version'] == 1
        
        # 验证数据已存储
        latest = self.memory_cache.get_latest_data(room_id)
        assert latest is not None
        assert latest['room_id'] == room_id
    
    @pytest.mark.asyncio
    async def test_realtime_data_retrieval(self):
        """测试实时数据获取"""
        room_id = "room_test_06"
        session_id = "session_006"
        
        # 先上传一些数据
        for i in range(5):
            data = self.data_generator.generate_eeg_data(room_id, session_id)
            await self.http_client.upload_data(data)
            await asyncio.sleep(0.01)
        
        # 获取实时数据
        response = await self.http_client.get_realtime_data(room_id)
        
        assert response['room_id'] == room_id
        assert response['version'] == 5
        assert response['data']['latest'] is not None
        assert 'recent' in response['data']
    
    @pytest.mark.asyncio
    async def test_version_checking(self):
        """测试版本检查机制"""
        room_id = "room_test_07"
        session_id = "session_007"
        
        # 初始版本检查
        version_resp = await self.http_client.check_version(room_id)
        assert version_resp['version'] == 0
        assert not version_resp['has_new_data']
        
        # 上传数据后检查
        data = self.data_generator.generate_eeg_data(room_id, session_id)
        await self.http_client.upload_data(data)
        
        version_resp = await self.http_client.check_version(room_id)
        assert version_resp['version'] == 1
        assert version_resp['has_new_data']
    
    @pytest.mark.asyncio
    async def test_client_polling_simulation(self):
        """测试客户端轮询模拟"""
        room_id = "room_test_08"
        session_id = "session_008"
        
        # 模拟客户端状态
        client_version = 0
        client_timestamp = time.time()
        
        # 模拟数据上传（APK端）
        async def data_uploader():
            for i in range(10):
                data = self.data_generator.generate_eeg_data(room_id, session_id)
                await self.http_client.upload_data(data)
                await asyncio.sleep(0.1)  # 每100ms上传一次
        
        # 模拟数据消费（其他服务）
        received_updates = []
        
        async def data_consumer():
            nonlocal client_version, client_timestamp
            
            for _ in range(15):  # 轮询15次
                # 检查版本
                version_resp = await self.http_client.check_version(room_id)
                
                if version_resp['version'] > client_version:
                    # 获取新数据
                    data_resp = await self.http_client.get_realtime_data(
                        room_id, since=client_timestamp
                    )
                    
                    if data_resp['data']['recent']:
                        received_updates.extend(data_resp['data']['recent'])
                    
                    client_version = version_resp['version']
                    client_timestamp = time.time()
                
                await asyncio.sleep(0.08)  # 每80ms轮询一次
        
        # 并发运行上传和消费
        await asyncio.gather(
            data_uploader(),
            data_consumer()
        )
        
        # 验证结果
        assert len(received_updates) > 0
        assert client_version > 0
    
    def test_performance_metrics(self):
        """测试性能指标"""
        room_id = "room_test_09"
        
        # 执行一系列操作
        for i in range(100):
            data = self.data_generator.generate_eeg_data(room_id, "session_009")
            self.memory_cache.add_data(room_id, data)
            
            # 随机读取
            if random.random() > 0.5:
                self.memory_cache.get_latest_data(room_id)
        
        # 获取统计信息
        stats = self.memory_cache.get_cache_stats()
        
        assert stats['total_rooms'] >= 1
        assert stats['data_updates'] == 100
        assert stats['total_requests'] > 0
        assert 0 <= stats['cache_hit_rate'] <= 1
    
    def test_memory_limits(self):
        """测试内存限制"""
        # 创建小容量缓存
        limited_cache = MemoryCache(max_rooms=3, max_window_size=5)
        
        # 添加超过限制的房间数据
        for i in range(5):
            room_id = f"room_limit_{i}"
            data = self.data_generator.generate_eeg_data(room_id, "session_limit")
            limited_cache.add_data(room_id, data)
        
        # 应该只保留最新的3个房间
        assert len(limited_cache.latest_data) <= 3
        
        # 测试时间窗口限制
        room_id = "room_window_test"
        for i in range(10):
            data = self.data_generator.generate_eeg_data(room_id, "session_window")
            limited_cache.add_data(room_id, data)
        
        # 时间窗口应该被限制
        assert len(limited_cache.time_window_data[room_id]) <= 5


# 性能测试
class TestPerformance:
    """性能测试"""
    
    def setup_method(self):
        self.memory_cache = MemoryCache()
        self.data_generator = EEGDataGenerator()
    
    def test_high_frequency_updates(self):
        """测试高频更新性能"""
        room_id = "perf_test_room"
        start_time = time.time()
        
        # 高频更新
        for i in range(1000):
            data = self.data_generator.generate_eeg_data(room_id, "perf_session")
            self.memory_cache.add_data(room_id, data)
        
        end_time = time.time()
        duration = end_time - start_time
        
        print(f"1000次更新耗时: {duration:.3f}秒")
        print(f"平均每次更新: {duration/1000*1000:.3f}毫秒")
        
        # 性能要求：每次更新应该在1ms以内
        assert duration / 1000 < 0.001
    
    def test_concurrent_performance(self):
        """测试并发性能"""
        results = {'total_operations': 0, 'start_time': time.time()}
        
        def worker(worker_id, operations=200):
            room_id = f"perf_room_{worker_id}"
            for i in range(operations):
                data = self.data_generator.generate_eeg_data(room_id, f"perf_session_{worker_id}")
                self.memory_cache.add_data(room_id, data)
                
                # 随机读取
                if i % 10 == 0:
                    self.memory_cache.get_latest_data(room_id)
                
                results['total_operations'] += 1
        
        # 启动多个并发线程
        threads = []
        for i in range(10):
            thread = threading.Thread(target=worker, args=(i,))
            threads.append(thread)
            thread.start()
        
        for thread in threads:
            thread.join()
        
        end_time = time.time()
        duration = end_time - results['start_time']
        
        print(f"并发测试: {results['total_operations']}次操作耗时: {duration:.3f}秒")
        print(f"平均TPS: {results['total_operations']/duration:.1f}")
        
        # 性能要求：TPS应该大于1000
        assert results['total_operations'] / duration > 1000


if __name__ == "__main__":
    # 运行基本测试
    print("开始TCP迁移综合测试...")
    
    # 创建测试实例
    test_instance = TestTCPMigration()
    test_instance.setup_method()
    
    # 运行同步测试
    print("\n1. 测试内存缓存基本操作...")
    test_instance.test_memory_cache_basic_operations()
    print("✓ 通过")
    
    print("\n2. 测试时间窗口数据...")
    test_instance.test_time_window_data()
    print("✓ 通过")
    
    print("\n3. 测试缓存过期和清理...")
    test_instance.test_cache_expiry_and_cleanup()
    print("✓ 通过")
    
    print("\n4. 测试并发访问安全性...")
    test_instance.test_concurrent_access()
    print("✓ 通过")
    
    print("\n5. 测试性能指标...")
    test_instance.test_performance_metrics()
    print("✓ 通过")
    
    print("\n6. 测试内存限制...")
    test_instance.test_memory_limits()
    print("✓ 通过")
    
    # 运行异步测试
    async def run_async_tests():
        print("\n7. 测试HTTP上传流程...")
        await test_instance.test_http_upload_flow()
        print("✓ 通过")
        
        print("\n8. 测试实时数据获取...")
        await test_instance.test_realtime_data_retrieval()
        print("✓ 通过")
        
        print("\n9. 测试版本检查机制...")
        await test_instance.test_version_checking()
        print("✓ 通过")
        
        print("\n10. 测试客户端轮询模拟...")
        await test_instance.test_client_polling_simulation()
        print("✓ 通过")
    
    # 运行异步测试
    asyncio.run(run_async_tests())
    
    # 运行性能测试
    print("\n=== 性能测试 ===")
    perf_test = TestPerformance()
    perf_test.setup_method()
    
    print("\n11. 测试高频更新性能...")
    perf_test.test_high_frequency_updates()
    print("✓ 通过")
    
    print("\n12. 测试并发性能...")
    perf_test.test_concurrent_performance()
    print("✓ 通过")
    
    print("\n🎉 所有测试通过！TCP迁移方案验证成功！")
    
    # 显示最终统计
    stats = test_instance.memory_cache.get_cache_stats()
    print(f"\n=== 最终统计 ===")
    print(f"总房间数: {stats['total_rooms']}")
    print(f"总请求数: {stats['total_requests']}")
    print(f"缓存命中率: {stats['cache_hit_rate']:.2%}")
    print(f"数据更新次数: {stats['data_updates']}")
    print(f"清理运行次数: {stats['cleanup_runs']}")